import torch
import torchvision
from torch.utils.data import DataLoader

dataset = torchvision.datasets.CIFAR10(root='./data/CIFAR10', train=False, download=True,
                                       transform=torchvision.transforms.ToTensor())
dataloder = DataLoader(dataset, batch_size=64)

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(196608, 10)

    def forward(self, x):
        x = self.fc1(x)
        return x

net = Net()

for step, (image, label) in enumerate(dataloder):
    output = net(torch.flatten(image))
    print(output.shape)